using System;
using System.Collections.Generic;
using Latios.Authoring;
using Unity.Collections;
using Unity.Entities;
using Unity.Mathematics;
using UnityEngine;

namespace Latios.Unika.Authoring
{
    /// <summary>
    /// A base class for authoring multiple scripts at once. This is an advanced API, so please pay close attention.
    /// </summary>
    public abstract class UnikaMultiScriptAuthoring : MonoBehaviour, IUnikaAuthoringScriptCounter
    {
        private protected static List<IUnikaAuthoringScriptCounter> m_scriptsCache = new List<IUnikaAuthoringScriptCounter>();

        /// <summary>
        /// Implement this method to determine if and how many scripts this authoring component will produce.
        /// Return -1 to completely cancel baking. Return 0 to not bake scripts but still bake other data.
        /// This method will be called frequently from various points in baking. It is crucial that this method
        /// returns -1 if the script baking will be canceled for any reason, as otherwise ScriptRefs obtained
        /// from any baked scripts on the target entity at bake time may be corrupted.
        /// </summary>
        /// <returns>A positive number for the exact number of scripts this authoring component will bake. -1 to cancel baking.</returns>
        public abstract int GetScriptCountToGenerate();

        /// <summary>
        /// Implement this method to control the baking of all the scripts.
        /// </summary>
        /// <param name="baker">A baker used to generate the script</param>
        /// <param name="toAssign">A span of objects to assign the script data to, as well as configure the user byte and flags</param>
        /// <param name="smartPostProcessTargets">A span of baking-only entities from which you can receive a reference to the script in an ISmartPostProcessItem
        /// before the script is merged into the scripts buffer. Each index corresponds to the same index in toAssign.</param>
        public abstract void Bake(IBaker baker, Span<AuthoredScriptAssignment> toAssign, Span<Entity> smartPostProcessTargets);

        /// <summary>
        /// Gets the nth untyped ScriptRef generated by this authoring component. Call this from other authoring scripts or bakers.
        /// </summary>
        /// <param name="baker">The baker being used to bake whatever consumes the ScriptRef</param>
        /// <param name="nthScriptIndex">The nth script index this authoring component generates, from 0 to the returned value of GetScriptCountToGenerate() - 1</param>
        /// <param name="transformUsageFlags">The transform flags that should be added to the script's entity, if any</param>
        /// <returns>An untyped ScriptRef for the script generated by this authoring component, or a Null ScriptRef if this authoring is not valid.</returns>
        public ScriptRef GetNthScriptRef(IBaker baker, int nthScriptIndex, TransformUsageFlags transformUsageFlags = TransformUsageFlags.None)
        {
            if (!enabled)
                return default;

            m_scriptsCache.Clear();
            baker.GetComponents(this, m_scriptsCache);
            for (int i = 0, validBefore = 0; i < m_scriptsCache.Count; i++)
            {
                var s = m_scriptsCache[i];
                if (ReferenceEquals(s, this))
                {
                    var scriptsToGenerate = GetScriptCountToGenerate();
                    if (scriptsToGenerate > 0 && scriptsToGenerate > nthScriptIndex)
                    {
                        return new ScriptRef
                        {
                            m_entity            = baker.GetEntity(this, transformUsageFlags),
                            m_instanceId        = validBefore + 1 + nthScriptIndex,
                            m_cachedHeaderIndex = validBefore + 1 + nthScriptIndex
                        };
                    }
                    else
                        return default;
                }
                else
                    validBefore += s.CountScripts();
            }
            return default;
        }

        /// <summary>
        /// Assigns the Script and extracts a typed ScriptRef which can be used to assign references to other components being baked,
        /// or even other scripts
        /// </summary>
        /// <typeparam name="T">The type of script to be generated</typeparam>
        /// <param name="baker">The baker being used to bake this script</param>
        /// <param name="script">The script field data the script should be initialized with</param>
        /// <param name="assignments">The span of assignment objects to assign the script data to</param>
        /// <param name="nthScriptIndex">The index in the span to assign this particular script data to and extract a typed ScriptRef for</param>
        /// <param name="storedScript">A NativeArray of length 1 referencing the script assigned to the nth assignment, which can be used for further ScriptRef patching</param>
        /// <param name="transformUsageFlags">The transform flags that should be added to the script's entity, if any</param>
        /// <returns>A typed ScriptRef for the assigned script</returns>
        protected ScriptRef<T> AssignNthAndGetTypedRef<T>(IBaker baker,
                                                          ref T script,
                                                          Span<AuthoredScriptAssignment> assignments,
                                                          int nthScriptIndex,
                                                          out NativeArray<T>             storedScript,
                                                          TransformUsageFlags transformUsageFlags = TransformUsageFlags.None)
            where T : unmanaged, IUnikaScript, IUnikaScriptGen
        {
            assignments[nthScriptIndex].Assign(ref script);
            var scriptRef = GetNthScriptRef(baker, nthScriptIndex, transformUsageFlags);
            storedScript  = assignments[nthScriptIndex].scriptPayload.Reinterpret<T>(1);
            return new ScriptRef<T>
            {
                m_entity            = scriptRef.m_entity,
                m_instanceId        = scriptRef.m_instanceId,
                m_cachedHeaderIndex = scriptRef.m_cachedHeaderIndex,
            };
        }

        int IUnikaAuthoringScriptCounter.CountScripts()
        {
            return enabled ? GetScriptCountToGenerate() : 0;
        }
    }

    [BakeDerivedTypes]
    public class UnikaMultiScriptAuthoringBaker : Baker<UnikaMultiScriptAuthoring>
    {
        public override void Bake(UnikaMultiScriptAuthoring authoring)
        {
            var scriptCount = authoring.GetScriptCountToGenerate();
            if (scriptCount < 0)
                return;

            if (scriptCount == 0)
            {
                authoring.Bake(this, default, default);
                return;
            }

            var entities = new NativeArray<Entity>(scriptCount, Allocator.Temp, NativeArrayOptions.UninitializedMemory);
            CreateAdditionalEntities(entities, TransformUsageFlags.None, true);
            Span<AuthoredScriptAssignment> assignments = stackalloc AuthoredScriptAssignment[scriptCount];
            authoring.Bake(this, assignments, entities.AsSpan());
            int assignmentIndex = 0;
            var flags           = TransformUsageFlags.None;
            foreach (var assignment in assignments)
            {
                if (assignment.scriptPayload.Length == 0)
                    throw new System.InvalidOperationException(
                        $"The AuthoredScriptAssignment was left unassigned for object {authoring.name} at index {assignmentIndex}. Please fix this issue before attempting to enter play mode, or else Unity may crash.");
                assignmentIndex++;
                flags |= assignment.transformUsageFlags;
            }

            GetEntity(flags);
            assignmentIndex = 0;
            foreach (var assignment in assignments)
            {
                var entity = entities[assignmentIndex];
                var bytes  = AddBuffer<BakedScriptByte>(entity).Reinterpret<byte>();
                bytes.AddRange(assignment.scriptPayload);
                AddComponent(entity, new BakedScriptMetadata
                {
                    scriptRef  = authoring.GetNthScriptRef(this, assignmentIndex),
                    userFlagA  = assignment.userFlagA,
                    scriptType = assignment.scriptType,
                    userByte   = assignment.userByte,
                    userFlagB  = assignment.userFlagB,
                });
                assignmentIndex++;
            }
        }
    }
}

